package com.ruoyi.common.utils;

import java.awt.image.BufferedImage;
import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.IOException;
import java.io.InputStream;
import java.net.URL;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Random;

import javax.imageio.ImageIO;

import org.datavec.api.io.labels.ParentPathLabelGenerator;
import org.datavec.api.split.FileSplit;
import org.datavec.image.loader.NativeImageLoader;
import org.datavec.image.recordreader.ImageRecordReader;
import org.deeplearning4j.api.storage.StatsStorage;
import org.deeplearning4j.datasets.datavec.RecordReaderDataSetIterator;
import org.deeplearning4j.earlystopping.EarlyStoppingConfiguration;
import org.deeplearning4j.earlystopping.EarlyStoppingResult;
import org.deeplearning4j.earlystopping.saver.LocalFileModelSaver;
import org.deeplearning4j.earlystopping.scorecalc.ClassificationScoreCalculator;
import org.deeplearning4j.earlystopping.termination.MaxEpochsTerminationCondition;
import org.deeplearning4j.earlystopping.trainer.EarlyStoppingTrainer;
import org.deeplearning4j.nn.multilayer.MultiLayerNetwork;
import org.deeplearning4j.optimize.listeners.ScoreIterationListener;
import org.deeplearning4j.ui.api.UIServer;
import org.deeplearning4j.ui.stats.StatsListener;
import org.deeplearning4j.ui.storage.InMemoryStatsStorage;
import org.deeplearning4j.util.ModelSerializer;
import org.deeplearning4j.zoo.ZooModel;
import org.deeplearning4j.zoo.model.LeNet;
import org.nd4j.evaluation.classification.Evaluation;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.nd4j.linalg.dataset.api.iterator.DataSetIterator;
import org.nd4j.linalg.dataset.api.preprocessor.DataNormalization;
import org.nd4j.linalg.dataset.api.preprocessor.ImagePreProcessingScaler;
import org.nd4j.linalg.learning.config.AdaDelta;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;

import com.ruoyi.project.system.picture.domain.Picture;
import com.ruoyi.project.system.picture.service.IPictureService;

public class DL4JUtils {
	public static String server = "http://localhost/";
	private static Logger log = LoggerFactory.getLogger(DL4JUtils.class);
	/**图片信息*/
 	public static int height = 32;
 	public static int width = 32;
 	public static int channels = 3;
 	/**神经网络参数*/
 	public static int rngseed = 123;//随机数种子
 	public static Random randNumGen = new Random(rngseed);//用于文件随机分割
 	public static int batchSize = 128;
 	public static int outputNum = 8;//分类数目
 	public static int numEpochs = 10;//迭代次数
    /**用于处理图片预测结果*/
 	public static HashMap<Integer, String> rsmap = new HashMap<Integer, String>();
 	public static String[] cnName = {"飞机","小车","鹿","小狗","蛙","马","船","卡车"};
    /**模型*/
 	public static MultiLayerNetwork model=null;
	public static void trainModel(String trainPath,String testPath) throws IOException {
		ParentPathLabelGenerator labelMaker = new ParentPathLabelGenerator();//为训练集创建标签，标签名为父文件夹名
		DataNormalization scaler = new ImagePreProcessingScaler(0,1);//归一化处理器
		//预处理训练集:
        File trainData = new File(trainPath);
        FileSplit train = new FileSplit(trainData, NativeImageLoader.ALLOWED_FORMATS,randNumGen);
        ImageRecordReader recordReader_train = new ImageRecordReader(height,width,channels,labelMaker);
        recordReader_train.initialize(train);
        DataSetIterator mnistTrain = new RecordReaderDataSetIterator(recordReader_train,batchSize,1,outputNum);
        scaler.fit(mnistTrain);
        mnistTrain.setPreProcessor(scaler);
      //预处理测试集:
        File testData = new File(testPath);
        FileSplit test = new FileSplit(testData,NativeImageLoader.ALLOWED_FORMATS,randNumGen);
        ImageRecordReader recordReader_test = new ImageRecordReader(height,width,channels,labelMaker);
        recordReader_test.initialize(test);
        DataSetIterator mnistTest = new RecordReaderDataSetIterator(recordReader_test,batchSize,1,outputNum);
        scaler.fit(mnistTest);
        mnistTest.setPreProcessor(scaler);	
        log.info("Build model....");
        /**使用DL4J提供的LeNet神经网络*/
        ZooModel zoomodel=LeNet.builder().updater(new AdaDelta()).seed(rngseed).inputShape(new int[] {channels, height, width}).numClasses(outputNum).build();
        /**网络初始化*/
        MultiLayerNetwork model = zoomodel.init();
        /**实例化训练用户界面，用于监测训练情况*/
        UIServer uiServer =UIServer.getInstance();
        StatsStorage statsStorage = new InMemoryStatsStorage(); 
        uiServer.attach(statsStorage);
        /**设置得分和状态监听器*/
        model.setListeners(new ScoreIterationListener(10),new StatsListener(statsStorage));
        
        /**配置早停法*/
        EarlyStoppingConfiguration esConf = new EarlyStoppingConfiguration.Builder()
                .epochTerminationConditions(new MaxEpochsTerminationCondition(10))//最多迭代10次
                .scoreCalculator(new ClassificationScoreCalculator(Evaluation.Metric.ACCURACY, mnistTest))//分类任务得分计算器
                .evaluateEveryNEpochs(1)//每迭代一次进行计算
                .modelSaver(new LocalFileModelSaver(new File("C:" + "/lenet-model-es.zip")))//保存得分最高的模型
                .build();

        EarlyStoppingTrainer trainer = new EarlyStoppingTrainer(esConf,model,mnistTrain);
        
        log.info("Train model....");
        /**执行早停训练*/
        EarlyStoppingResult result = trainer.fit();

        log.info("Evaluate model....");
        /**评估模型训练效果*/
        Evaluation eval = model.evaluate(mnistTest);
        log.info(eval.stats());
        log.info("****************finished********************");
	}
	public static void loadModel(String modelPath) throws IOException {
		for(int i = 0;i<outputNum;i++) {
        	rsmap.put(i, cnName[i]);
        }
        //恢复模型
		model =ModelSerializer.restoreMultiLayerNetwork(new File(modelPath));
	}
	public static List<Picture> useModelPredict(List<Picture> pictures) throws IOException {
		NativeImageLoader imageLoader = new NativeImageLoader(height,width,channels);//实例化图片加载器
		DataNormalization scaler = new ImagePreProcessingScaler(0, 1);//归一化处理器
		INDArray imgArray;
		List<Picture> rs = new ArrayList<Picture>();
		for(Picture picture:pictures) {
			BufferedImage image = ImageIO.read(new URL(server+picture.getSrcPath()));//使用ImageIO读取图片数据
			ByteArrayOutputStream os = new ByteArrayOutputStream();
			ImageIO.write(image, "jpg", os);//将图片数据写入到ByteArrayOutputStream
		    InputStream input = new ByteArrayInputStream(os.toByteArray());//将ByteArrayOutputStream转换为InputStream
		    imgArray = imageLoader.asMatrix(input);//将InputStream转换为ND4j矩阵数组
		    scaler.transform(imgArray);//归一化处理
		    picture.setRemark(rsmap.get(model.predict(imgArray)[0]));//使用模型预测图片并将识别结果保存到数据库
		    rs.add(picture);
		}
		return rs;
	}
}
